#pragma once

#ifdef _WIN32

#include "Win32MiniDump.h"
#include <windows.h>
#include <stdio.h>
#include <assert.h>
#include <time.h>
#include <stdlib.h>
#include <dbghelp.h>
#include <tchar.h>
#include <strsafe.h>
#include "base/core/log.h"

Win32MiniDump* Win32MiniDump::s_pMiniDumper = NULL;
LPCRITICAL_SECTION Win32MiniDump::s_pCriticalSection = NULL;

typedef BOOL(WINAPI *MINIDUMPWRITEDUMP)(HANDLE hProcess,
										DWORD dwPid,
										HANDLE hFile,
										MINIDUMP_TYPE DumpType,
										CONST PMINIDUMP_EXCEPTION_INFORMATION ExceptionParam,
										CONST PMINIDUMP_USER_STREAM_INFORMATION UserStreamParam,
										CONST PMINIDUMP_CALLBACK_INFORMATION CallbackParam);

Win32MiniDump::Win32MiniDump()
{
	assert(!s_pMiniDumper);
	s_pMiniDumper = this;

	::SetUnhandledExceptionFilter(UnhandledExceptionHandler);

	s_pCriticalSection = new CRITICAL_SECTION;
	if (s_pCriticalSection)
	{ 
		InitializeCriticalSection(s_pCriticalSection);
	}
}

Win32MiniDump::~Win32MiniDump(void)
{
	if (s_pCriticalSection)
	{
		DeleteCriticalSection(s_pCriticalSection);
		delete s_pCriticalSection;
	}
}

long Win32MiniDump::UnhandledExceptionHandler(_EXCEPTION_POINTERS *pExceptionInfo)
{
	if (!s_pMiniDumper)
	{ 
		return EXCEPTION_CONTINUE_SEARCH;
	}
	return s_pMiniDumper->WriteMiniDump(pExceptionInfo);
}

void Win32MiniDump::SetMiniDumpFileName(void)
{
	time_t currentTime;
	time(&currentTime);
	wsprintf(m_szMiniDumpPath, "%s%s.%ld.dmp", m_szAppPath, m_szAppBaseName, currentTime);
}

/*
Desc: The method acts as a potential workaround for the fact that the
	  current thread may not have a token assigned to it, and if not, the
	  process token is received.
*/
bool Win32MiniDump::CheckThreadToken(HANDLE* phToken)
{
	*phToken = NULL;

	if (!OpenThreadToken(GetCurrentThread(), TOKEN_QUERY | TOKEN_ADJUST_PRIVILEGES, TRUE, phToken))
	{
		if (GetLastError() == ERROR_NO_TOKEN)
		{
			if (!OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY | TOKEN_ADJUST_PRIVILEGES, phToken))
			{ 
				return false;
			}
		}
		else
		{ 
			return false;
		}
	}
	return true;
}

/*
Desc: Since a MiniDump contains a lot of meta-data about the OS and
	  application state at the time of the dump, it is a rather privileged
	  operation. This means we need to set the SeDebugPrivilege to be able
	  to call MiniDumpWriteDump.
*/
BOOL Win32MiniDump::EnablePrivilege(LPCTSTR pszPriv, HANDLE hToken, TOKEN_PRIVILEGES* ptpOld)
{
	BOOL bOk = FALSE;

	TOKEN_PRIVILEGES tp;
	tp.PrivilegeCount = 1;
	tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
	bOk = LookupPrivilegeValue(0, pszPriv, &tp.Privileges[0].Luid);

	if (bOk)
	{
		DWORD cbOld = sizeof(*ptpOld);
		bOk = AdjustTokenPrivileges(hToken, FALSE, &tp, cbOld, ptpOld, &cbOld);
	}

	return (bOk && (ERROR_NOT_ALL_ASSIGNED != GetLastError()));
}

BOOL Win32MiniDump::RestorePrivilege(HANDLE hToken, TOKEN_PRIVILEGES* ptpOld)
{
	BOOL bOk = AdjustTokenPrivileges(hToken, FALSE, ptpOld, 0, NULL, NULL);
	return (bOk && (ERROR_NOT_ALL_ASSIGNED != GetLastError()));
}

long Win32MiniDump::WriteMiniDump(_EXCEPTION_POINTERS *pExceptionInfo)
{
	long retval = EXCEPTION_CONTINUE_SEARCH;
	m_pExceptionInfo = pExceptionInfo;

	HANDLE hThreadToken = NULL;
	if (!CheckThreadToken(&hThreadToken))
	{
		return FALSE;
	}

	// You have to find the right dbghelp.dll.   
	HMODULE hDll = NULL;
	TCHAR szDbgHelpPath[MAX_PATH];
	if (GetModuleFileName(NULL, m_szAppPath, _MAX_PATH))
	{
		TCHAR *pSlash = strrchr(m_szAppPath, '\\');
		if (pSlash)
		{
			_tcscpy(m_szAppBaseName, pSlash + 1);
			*(pSlash + 1) = 0;
		}
		strcpy(szDbgHelpPath, m_szAppPath);
		strcat(szDbgHelpPath, "DBGHELP.DLL");
		hDll = ::LoadLibrary(szDbgHelpPath);
	}

	if (hDll == NULL)
	{
		hDll = ::LoadLibrary("DBGHELP.DLL");
	}

	if (hDll)
	{
		// Get the address of the MiniDumpWriteDump function, which writes
		// user-mode mini-dump information to a specified file.
		MINIDUMPWRITEDUMP MiniDumpWriteDump = (MINIDUMPWRITEDUMP)::GetProcAddress(hDll, "MiniDumpWriteDump");
		if (MiniDumpWriteDump != NULL)
		{
			SetMiniDumpFileName();

			// Create the mini-dump file...  
			HANDLE hFile = ::CreateFile(m_szMiniDumpPath, GENERIC_WRITE, FILE_SHARE_WRITE, NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
			if (hFile != INVALID_HANDLE_VALUE)
			{
				_MINIDUMP_EXCEPTION_INFORMATION ExInfo;
				ExInfo.ThreadId = ::GetCurrentThreadId();
				ExInfo.ExceptionPointers = pExceptionInfo;
				ExInfo.ClientPointers = NULL;

				TOKEN_PRIVILEGES tp;
				BOOL bPrivilegeEnabled = EnablePrivilege(SE_DEBUG_NAME, hThreadToken, &tp);
				BOOL bOk;

				// DBGHELP.dll is not thread-safe, so we need to restrict access...  
				EnterCriticalSection(s_pCriticalSection);
				{
					// Write out the mini-dump data to the file...  
					bOk = MiniDumpWriteDump(GetCurrentProcess(), GetCurrentProcessId(), hFile, MiniDumpNormal, &ExInfo, NULL, NULL);
				}
				LeaveCriticalSection(s_pCriticalSection);

				// Restore the privileges when done
				if (bPrivilegeEnabled)
				{ 
					RestorePrivilege(hThreadToken, &tp);
				}

				if (bOk)
				{
					retval = EXCEPTION_EXECUTE_HANDLER;
				}
				::CloseHandle(hFile);

				LogErrPrint("[Win32MiniDump] create Win32MiniDump succ");
			}
		}
	}
	else
	{
		LogErrPrint("[Win32MiniDump] can not load DBGHELP.DLL");
	}

	TerminateProcess(GetCurrentProcess(), 0);
	return retval;
}

#endif